from cython.parallel import prange
cimport cython
import numpy as np
cimport numpy as np
from libc.math cimport sqrt, NAN  # 导入 NAN

@cython.boundscheck(False)
@cython.wraparound(False)
def ts_argmax(double[:, :] x, int window=10):
    cdef int n = x.shape[0]
    cdef int m = x.shape[1]
    cdef double[:, :] result = np.empty((n, m), dtype=np.float64)
    cdef int i, j, k
    cdef double max_val, val
    cdef int argmax_idx

    with nogil:
        for j in prange(m, schedule='dynamic'):
            for i in range(n):
                if i >= window - 1:
                    max_val = x[i - window + 1, j]
                    argmax_idx = 0
                    for k in range(1, window):
                        val = x[i - window + 1 + k, j]
                        if val > max_val:
                            max_val = val
                            argmax_idx = k
                    result[i, j] = argmax_idx + 1
                else:
                    result[i, j] = NAN
    return np.array(result)

@cython.boundscheck(False)
@cython.wraparound(False)
def stddev(double[:, :] x, int window=10):
    cdef int n = x.shape[0]
    cdef int m = x.shape[1]
    cdef double[:, :] result = np.empty((n, m), dtype=np.float64)
    cdef int j

    with nogil:
        for j in prange(m, schedule='dynamic'):
            compute_stddev_column(x[:, j], result[:, j], window)
    return np.array(result)

cdef void compute_stddev_column(double[:] x_col, double[:] result_col, int window) nogil:
    cdef int n = x_col.shape[0]
    cdef int i, k
    cdef double mean, sum_sq, val, delta

    for i in range(n):
        if i >= window - 1:
            mean = 0.0
            sum_sq = 0.0
            for k in range(window):
                val = x_col[i - window + 1 + k]
                mean += val
            mean /= window
            for k in range(window):
                val = x_col[i - window + 1 + k]
                delta = val - mean
                sum_sq += delta * delta
            result_col[i] = sqrt(sum_sq / (window - 1))
        else:
            result_col[i] = NAN

@cython.boundscheck(False)  # 关闭边界检查以提高性能
@cython.wraparound(False)   # 关闭负索引检查以提高性能
def rolling_std_optimized(np.ndarray[np.float64_t, ndim=2] x, int window=10):
    cdef int rows = x.shape[0]
    cdef int cols = x.shape[1]
    cdef int i, j

    # 结果矩阵，初始化为 NaN
    cdef np.ndarray[np.float64_t, ndim=2] result = np.full((rows, cols), np.nan, dtype=np.float64)

    # 临时变量来计算滚动窗口的和和平方和
    cdef double window_sum, window_sq_sum, mean, variance

    for j in range(cols):
        window_sum = 0.0
        window_sq_sum = 0.0

        # 初始化第一个窗口
        for i in range(window):
            window_sum += x[i, j]
            window_sq_sum += x[i, j]**2

        # 滚动窗口计算
        for i in range(window - 1, rows):
            mean = window_sum / window
            variance = (window_sq_sum / window) - mean**2
            result[i, j] = sqrt(variance)

            # 更新窗口和窗口平方和
            if i + 1 < rows:
                window_sum += x[i + 1, j] - x[i - window + 1, j]
                window_sq_sum += x[i + 1, j]**2 - x[i - window + 1, j]**2

    return result


def rankdata(double[:, :] a, method='average', int axis=1):
    cdef int n = a.shape[0]
    cdef int m = a.shape[1]
    cdef double[:, :] ranks = np.empty((n, m), dtype=np.float64)
    cdef int i

    if axis == 1:
        for i in range(n):
            ranks[i, :] = rank1d(a[i, :], method)
    else:
        # 如果需要沿其他轴排序，需要相应修改
        pass

    return np.array(ranks)

def rank1d(double[:] arr, method):
    cdef int size = arr.shape[0]
    cdef int[:] order = np.argsort(arr)
    cdef int[:] ranks = np.empty(size, dtype=np.int32)
    cdef int i

    for i in range(size):
        ranks[order[i]] = i + 1

    # 根据不同的方法进行处理，这里以 'min' 方法为例
    return ranks.astype(np.float64) / size